import threading
from socket import *
import time
import cv2
import numpy as np
from PIL import Image
from yolo import YOLO

yolo = YOLO()

cap = cv2.VideoCapture(0)
# 统一图像大小为640*480
cap.set(cv2.CAP_PROP_FRAME_WIDTH, 320)
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 240)

Host = 'localhost'
Port = 5555
sock = socket(AF_INET, SOCK_DGRAM)  # 创建UDP套接字
print("Sending...")


def send():
    while True:
        time.sleep(0.1)
        # 读取某一帧
        ref, frame = cap.read()
        if not ref:
            break
        # 格式转变，BGRtoRGB
        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        # 转变成Image
        frame = Image.fromarray(np.uint8(frame))
        # 进行检测
        frame = np.array(yolo.detect_image(frame))
        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

        ret, send_data = cv2.imencode('.jpg', frame, [cv2.IMWRITE_JPEG_QUALITY, 50])  # 将图片转为byte
        sock.sendto(send_data, (Host, Port))

        # RGBtoBGR满足opencv显示格式
        # frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
        # cv2.imshow("video", frame)
        c = cv2.waitKey(1) & 0xff

        if c == 27:
            cap.release()
            break

    print("Video Detection Done!")
    cap.release()
    cv2.destroyAllWindows()


t1 = threading.Thread(target=send)
t1.start()



